from airflow import DAG
from airflow.operators.bash import BashOperator
from airflow.utils.dates import days_ago
from datetime import timedelta

params = {
	'python3': '/usr/bin/python3',
	'work_root': '/opt/efso_work/',
	'geps_root': '/opt/data/raw/GRAPES_GEPS/',
	'pi_root': '/g1/COMMONDATA/OPER/NWPC/GRAPES_GEPS/Prod-grib/',
	'obs_root': '/data/raw/',
	'efso_image': '~/containers/singularity/efso/efso-1.0.sif',
	'ntasks': 1
}

args = {
	'owner': 'airflow'
}

with DAG(
	dag_id='grapes_geps_efso',
	default_args=args,
	user_defined_macros=params,
	schedule_interval='0 00,12 * * *',
	start_date=days_ago(2),
	dagrun_timeout=timedelta(minutes=60),
) as dag:
	get_geps= BashOperator(
		task_id='get_geps',
		bash_command='''
{% set init_time= execution_date.subtract(days=1).format("YYYYMMDDHH") %}
{% set anal_time= execution_date.format("YYYYMMDDHH") %}
{% set data_dir = geps_root + '/' + init_time  %}
rsync -avP guojx@10.40.140.19:{{pi_root}}/{{init_time}}/grib2/gef.gra.0*.{{init_time}}000.grb2 {{geps_root}}/{{init_time}}/
rsync -avP guojx@10.40.140.19:{{pi_root}}/{{init_time}}/grib2/gef.gra.0*.{{init_time}}024.grb2 {{geps_root}}/{{init_time}}/
rsync -avP guojx@10.40.140.19:{{pi_root}}/{{anal_time}}/grib2/gef.gra.0*.{{anal_time}}000.grb2 {{geps_root}}/{{anal_time}}/
rsync -avP guojx@10.40.140.19:{{pi_root}}/{{anal_time}}/grib2/gef.gra.0*.{{anal_time}}024.grb2 {{geps_root}}/{{anal_time}}/
'''
	)
	prepare = BashOperator(
		task_id='prepare',
		bash_command='''
{% set init_time= execution_date.subtract(days=1).format("YYYYMMDDHH") %}
{% set work_dir = work_root + '/' + init_time %}
test -d {{work_dir}} || mkdir -p {{work_dir}}
'''
	)

	write_namelist = BashOperator(
		task_id='write_namelist',
		bash_command='''
{% set init_time= execution_date.subtract(days=1).format("YYYYMMDDHH") %}
{% set work_dir = work_root + '/' + init_time %}
cd {{work_dir}} && {{python3}} /opt/efso/src/write_namelist.py --time {{init_time}} \
--use-raob --raob-root /data/raw/cimiss/UPAR_CHN_MUL_FTM \
--use-amdar --amdar-root /data/raw/cimiss/UPAR_ARD_G_MUL_MUT_TAB \
--grapes-root {{geps_root}} \
--fcst-hour 24
'''
	)

#--use-profiler --profiler-root /data/raw/profiler \

	run_efso = BashOperator(
		task_id='run_efso',
		bash_command='''
{% set init_time= execution_date.subtract(days=1).format("YYYYMMDDHH") %}
{% set work_dir = work_root + '/' + init_time %}
cd {{work_dir}} && source {{work_root}}/env && \
mpiexec -np {{ntasks}} /opt/efso/build/efso_driver.exe namelist
'''
	)

	write_db = BashOperator(
		task_id='write_db',
		bash_command='''
{% set init_time= execution_date.subtract(days=1).format("YYYYMMDDHH") %}
{% set work_dir = work_root + '/' + init_time %}
cd {{work_dir}} && {{python3}} /opt/efso/src/write_db.py --time {{init_time}} \
--input obs_impact.{{init_time}}.txt
'''
	)
	clean_data = BashOperator(
		task_id='clean_data',
		bash_command='''
find {{geps_root}} -mtime +7 -exec rm -rf {} \;
'''
	)

	get_geps >> prepare
	prepare >> write_namelist
	write_namelist >> run_efso
	run_efso >> write_db
	clean_data
